#ifndef __GPUSAMATCH_KERNEL_CU__
#define __GPUSAMATCH_KERNEL_CU__


#include "GPUSAMatch_kernel.h"
#include "Settings.h"
#include "cuda_header.h"
#include <mian_util.h>
#include <thrust/scan.h>
#include <thrust/reduce.h>

//////////////////////////////////// CUDA kenrels //////////////////////////////////////////

__global__
void extractSA_kernel(unsigned int* d_qid, unsigned int* d_length, 
			const st_SA* d_sa, const int numSA) {
	const int numTotalThread = NUM_TOTAL_THREAD;
	const int globalThreadOffset = GLOBAL_THREAD_OFFSET;
	st_SA sa;

	for(int i = globalThreadOffset; i < numSA; i += numTotalThread) {
		sa = d_sa[i];
		d_qid[i] = sa.x;
		d_length[i] = sa.z - sa.y + 1;
	}
}


__device__
int reduceInWarp(char* data, const int idInWarp) {
	if(idInWarp < 16) {
		data[idInWarp] += data[idInWarp + 16];
	}
	//16 elements now

	if(idInWarp < 8) {
		data[idInWarp] += data[idInWarp + 8];
	}
	//8 elements now

	if(idInWarp < 4) {
		data[idInWarp] += data[idInWarp + 4];
	}
	//4 elements now

	if(idInWarp < 2) {
		data[idInWarp] += data[idInWarp + 2];
	}
	//2 elements now

	if(idInWarp < 1) {
		data[0] += data[1];
	}
	//1 element now

	return data[0];
}

__device__
bool isExactMatched(const char* read, const char* ref, const int length, 
			const int idInWarp, char* buf) {
	const int nloop = ceil(length/(double)WARP_SIZE);
	int i;

	for(int passId = 0; passId < nloop; passId += 1) {
		buf[idInWarp] = 0;

		i = passId*WARP_SIZE + idInWarp;
		if(i < length) {
			buf[idInWarp] = (read[i] == ref[i]) ? (0) : (1);
		}

		if(reduceInWarp(buf, idInWarp) > 0) {
			return false;
		}
	}
	return true;
}



typedef bool (*MATCH_FUNC)(const char*, const char*, const int, const int, char*);


__device__
bool isOneMismatched(const char* read, const char* ref, const int length, 
			const int idInWarp, char* buf) {
	const int nloop = ceil(length/(double)WARP_SIZE);
	int i;
	int numMis = 0;

	for(int passId = 0; passId < nloop; passId += 1) {
		buf[idInWarp] = 0;

		i = passId*WARP_SIZE + idInWarp;
		if(i < length) {
			buf[idInWarp] = (read[i] == ref[i]) ? (0) : (1);
		}

		numMis += reduceInWarp(buf, idInWarp);
		if(numMis > 1) {
			return false;
		}
	}
	
	if(1 == numMis) {
		return true;
	} else {
		return false;
	}

}


__device__
bool _1misBackward(const char* read, const char* ref, const int readLength, 
			const int idInWarp, char* s_flag) {

	//the second half is exact match
	if(!isExactMatched(read + readLength/2, ref + readLength/2, readLength/2, 
				idInWarp, s_flag)) {
		return false;
	}

	//the first half has and only has one mismatch
	if(!isOneMismatched(read, ref, readLength/2, 
				idInWarp, s_flag))  {
		return false;
	}

	return true;
}


__device__
bool _1misForward(const char* read, const char* ref, const int readLength, 
			const int idInWarp, char* s_flag) {

	//the first half is exact match
	if(!isExactMatched(read, ref, readLength/2, idInWarp, s_flag)) {
		return false;
	}

	//the second half has and only has one mismatch
	if(!isOneMismatched(read + readLength/2, ref + readLength/2, readLength/2, 
				idInWarp, s_flag))  {
		return false;
	}

	return true;
}

__inline__ __device__
void loadRead(char* s_read, const char* d_read, 
		const int rid, const int readLength,
		const int idInWarp) {
	for(int i = idInWarp; i < readLength; i += WARP_SIZE) {
			s_read[i] = d_read[rid*readLength + i];
	}
}



//parallel strategy: one warp to handle one SA
__global__
void cuda_match_1misbackward_kernel(
		/*parameters from D_MATCH_RESULT*/
		char* d_isMatched,
		/*parameters from D_READSET*/
		const char* d_read, 
		const int numRead, const int readLength,
		/*parameters from D_REF*/
		const char* d_ref, const unsigned int numBase,
		/*parameters from D_POS_HELPER*/	
		const unsigned int numSA, const unsigned int numPos, 
		const unsigned int* d_rid, 
		const unsigned int* d_offset, 
		const unsigned int* d_length, 
		const unsigned int* d_pos ) {
	const int globalThreadOffset = GLOBAL_THREAD_OFFSET;
	const unsigned int warpId = globalThreadOffset/WARP_SIZE;
	const unsigned int warpIdInBlock = threadIdx.x/WARP_SIZE;
	const unsigned int idInWarp = globalThreadOffset%WARP_SIZE;
	const unsigned int numWarp = NUM_TOTAL_THREAD/WARP_SIZE;
	unsigned int rid, start, length, p;
	extern __shared__ unsigned char s_data[];
	char* s_flag = (char*)(s_data + sizeof(char)*warpIdInBlock*WARP_SIZE);
	char* s_read = (char*)(s_data + sizeof(char)*blockDim.x +
				 sizeof(char)*readLength*warpIdInBlock);
	const char* p_ref;

	for(unsigned int sid = warpId; sid < numSA; sid += numWarp) {
		//copy the read to the shared memory
		rid = d_rid[sid];
		loadRead(s_read, d_read, rid, readLength, idInWarp);

		//process each position in the SA
		start = d_offset[sid];
		length = d_length[sid];
		for(unsigned int k = start; k < start + length; k += 1) {
			p = d_pos[k];
			p_ref = d_ref + p;
		
			//now we are ready to compare s_read and p_ref
			bool isMatched = _1misBackward(s_read, p_ref, readLength, 
							idInWarp, s_flag);
			
			//set the result for this read
			if(idInWarp == 0) {
				d_isMatched[k] = (isMatched) ? (1) : (0);
			}
		}
	}
}


__global__
void cuda_match_1misforward_kernel(
		/*parameters from D_MATCH_RESULT*/
		char* d_isMatched,
		/*parameters from D_READSET*/
		const char* d_read, 
		const int numRead, const int readLength,
		/*parameters from D_REF*/
		const char* d_ref, const unsigned int numBase,
		/*parameters from D_POS_HELPER*/	
		const unsigned int numSA, const unsigned int numPos, 
		const unsigned int* d_rid, 
		const unsigned int* d_offset, 
		const unsigned int* d_length, 
		const unsigned int* d_pos ) {
	const int globalThreadOffset = GLOBAL_THREAD_OFFSET;
	const unsigned int warpId = globalThreadOffset/WARP_SIZE;
	const unsigned int warpIdInBlock = threadIdx.x/WARP_SIZE;
	const unsigned int idInWarp = globalThreadOffset%WARP_SIZE;
	const unsigned int numWarp = NUM_TOTAL_THREAD/WARP_SIZE;
	unsigned int rid, start, length, p;
	extern __shared__ unsigned char s_data[];
	char* s_flag = (char*)(s_data + sizeof(char)*warpIdInBlock*WARP_SIZE);
	char* s_read = (char*)(s_data + sizeof(char)*blockDim.x +
				 sizeof(char)*readLength*warpIdInBlock);
	const char* p_ref;

	for(unsigned int sid = warpId; sid < numSA; sid += numWarp) {
		//copy the read to the shared memory
		rid = d_rid[sid];
		loadRead(s_read, d_read, rid, readLength, idInWarp);

		//process each position in the SA
		start = d_offset[sid];
		length = d_length[sid];
		for(unsigned int k = start; k < start + length; k += 1) {
			p = d_pos[k];
			p_ref = d_ref + p;
		
			//now we are ready to compare s_read and p_ref
			bool isMatched = _1misForward(s_read, p_ref, readLength, 
							idInWarp, s_flag);
			
			//set the result for this read
			if(idInWarp == 0) {
				d_isMatched[k] = (isMatched) ? (1) : (0);
			}
		}
	}
}


//////////////////////////////////// CPU gold functions, only for result validaction ////////////

void gold_match(D_MATCH_RESULT &d_matchResult, D_READSET &d_readset, 
		D_REF &d_ref, D_POS_HELPER &d_poshelper, 
		const int numMis) {
	const int numRead = d_readset.numRead;
	const int readLength = d_readset.readLength;
	const unsigned int numSA = d_poshelper.numSA;
	const unsigned int numPos = d_poshelper.numPos;
	unsigned int* rid = new unsigned int[numSA];
	unsigned int* offset = new unsigned int[numSA];
	unsigned int* length = new unsigned int[numSA];
	unsigned int* pos = new unsigned int[numPos];
	FROMGPU(rid, d_poshelper.d_rid, sizeof(unsigned int)*numSA);
	FROMGPU(offset, d_poshelper.d_offset, sizeof(unsigned int)*numSA);
	FROMGPU(length, d_poshelper.d_length, sizeof(unsigned int)*numSA);
	FROMGPU(pos, d_poshelper.d_pos, sizeof(unsigned int)*numPos);
	char* read = new char[readLength*numRead];
	FROMGPU(read, d_readset.d_read, sizeof(char)*readLength*numRead);
	char* ref = new char[d_ref.numBase];
	FROMGPU(ref, d_ref.d_ref, sizeof(char)*(d_ref.numBase));
	

	unsigned int numAlignment = 0;
	for(int sid = 0; sid < numSA; sid++) {
		const int start = offset[sid];
		const char* p_read = read + sizeof(char)*readLength*rid[sid];

		for(int k = start; k < start + length[sid]; k++) {
			const int p = pos[k];
			const char* p_ref = ref + p;
			int m = 0;
			for(int i = 0; i < readLength; i++) {
				if(p_read[i] != p_ref[i]) {
					m += 1;
				}
			}
			if(m == numMis) {
				numAlignment += 1;
			}
		}
	}

	cout << "goldMatch::numAlignment = " << numAlignment << endl;

	delete[] rid;
	delete[] offset;
	delete[] length;
	delete[] pos;
	delete[] read;
	delete[] ref;
}

//////////////////////////////////// kenrel APIs ///////////////////////////


/**
strategy:
1. we make each warp handle one SA, which corresponds to a read and a few positions
2. each thread corresponds to an 8-bit word, to record the number of mismatch for that thread
3. finally, we perform the reduction to get the number of mismatches for each read
*/

void cuda_match(D_MATCH_RESULT &d_matchResult, D_READSET &d_readset, 
		D_REF &d_ref, D_POS_HELPER &d_poshelper, 
		const MATCH_CASE matchCase) {
	//gold_match(d_matchResult, d_readset, d_ref, d_poshelper, numMis);


	//prepare
	const int numBlock = 1024;
	const int numThread = 128;
	//shared memory
	const int sharedMemSize = 
		sizeof(char)*(d_readset.readLength)*(numThread/WARP_SIZE) + 	//for read based per warp
		sizeof(char)*numThread;						//for result indication
	cout << "sharedMemSize = " << sharedMemSize << endl;

	if(matchCase == MIS1_BACKWARD) {
		cuda_match_1misbackward_kernel<<<numBlock, numThread, sharedMemSize>>>(
			//parameters from D_MATCH_RESULT
			d_matchResult.d_isMatched, 
			//parameters from D_READSET
			d_readset.d_read, d_readset.numRead, d_readset.readLength,
			//parameters from D_REF
			d_ref.d_ref, d_ref.numBase,
			//parameters from D_POS_HELPER
			d_poshelper.numSA, d_poshelper.numPos, 
			d_poshelper.d_rid, d_poshelper.d_offset, 
			d_poshelper.d_length, d_poshelper.d_pos);
		cutilCheckMsg("cuda_match_1misbackward_kernel");
	} else if(matchCase == MIS1_FORWARD) {
		cuda_match_1misforward_kernel<<<numBlock, numThread, sharedMemSize>>>(
			//parameters from D_MATCH_RESULT
			d_matchResult.d_isMatched, 
			//parameters from D_READSET
			d_readset.d_read, d_readset.numRead, d_readset.readLength,
			//parameters from D_REF
			d_ref.d_ref, d_ref.numBase,
			//parameters from D_POS_HELPER
			d_poshelper.numSA, d_poshelper.numPos, 
			d_poshelper.d_rid, d_poshelper.d_offset, 
			d_poshelper.d_length, d_poshelper.d_pos);
		cutilCheckMsg("cuda_match_1misforward_kernel");

	}else {
		ERROR_EXIT("NOT SUPPORTED YET!");
	}

	//report result
	char* h_isMatched = new char[d_poshelper.numPos];
	cout << "numPos = " << d_poshelper.numPos << endl;
	FROMGPU(h_isMatched, d_matchResult.d_isMatched, sizeof(char)*(d_poshelper.numPos));
	unsigned int numAlignment = 0;
	for(unsigned int i = 0; i < d_poshelper.numPos; i++) {
		numAlignment += h_isMatched[i];
	}
	cout << "cuda_match::numAlignment = " << numAlignment << endl;

	delete[] h_isMatched;

}



int cuda_extractSA(unsigned int* d_qid, unsigned int* d_offset,
                unsigned int* d_length, 
		const st_SA* d_sa, const int numSA) {

	extractSA_kernel<<<1024, 256>>>(d_qid, d_length, d_sa, numSA);
	cutilCheckMsg("extractSA_kernel");

	thrust::device_ptr<unsigned int> dptr_offset(d_offset);
	thrust::device_ptr<unsigned int> dptr_length(d_length);
	thrust::exclusive_scan(dptr_length, dptr_length + numSA, dptr_offset);
	
	return thrust::reduce(dptr_length, dptr_length + numSA);
}

#endif /*__GPUSAMATCH_KERNEL_CU__*/

